#!/usr/bin/env python

"""
This is a test of the database of contacts. It expects the db to be
in such state that
a) every pdb code corresponds to a set of hydrogen bonds and a set
   of near bonds
b) there are no other contacts than H2O-protein, H2O-DNA or DNA-protein
   (U bases are OK)
c) every contacts table corresponds to an atoms table
"""

__version__ = "$Id$"

import unittest
import sqlite3
import sys

class DatabaseTest(unittest.TestCase):
    def setUp(self, *args):
    	self.conn = sqlite3.connect(DB)
    	self.cur = self.conn.cursor()
    	
    	self.tables = self.cur.execute("""SELECT name FROM sqlite_master
                    WHERE type='table' AND name LIKE 'contacts_%'
                    ORDER BY name""").fetchall()
                    
        
    def testA(self):
        for (table,) in self.tables:
            q = self.cur.execute("""SELECT type, COUNT(*) FROM %s
                    GROUP BY type
                    ORDER BY type""" % table)
            res = q.fetchall()
            #print table, res
            self.failUnless(len(res) == 2) # there must be 2 types
            self.failUnless(res[0][1] > 0 and res[0][1] > 0) # which have nonzero number of records
            if res[0][1] > res[0][1]:
                print "Warning, in table %s is more H bonds than N bonds" % table
    
    def testB(self):
        cur2 = self.conn.cursor()
        for (table,) in self.tables:
            # there should be only contacts of DNA-* or HOH-*
            SQL = """SELECT donor.amino, acceptor.amino FROM
                     contacts_%(pdb)s AS contacts
                        JOIN atoms_%(pdb)s AS donor
                           ON donor.id = contacts.donor_id
                        JOIN atoms_%(pdb)s AS acceptor
                           ON acceptor.id = contacts.acceptor_id
                     WHERE donor.amino != "HOH" AND acceptor.amino != "HOH"
                       AND NOT ((donor.amino = "A" or
                                donor.amino = "G" or
                                donor.amino = "C" or
                                donor.amino = "U" or
                                donor.amino = "T") or
                                (acceptor.amino = "A" or
                                acceptor.amino = "G" or
                                acceptor.amino = "C" or
                                acceptor.amino = "U" or
                                acceptor.amino = "T"))
                     """
            pdb = table.split('_')[1]
            r = cur2.execute(SQL % {'pdb' : pdb})
            r = r.fetchall()
            self.failUnless(len(r) == 0, "table %s:\n%s" % (table, r))
            
            # we don't want HOH-HOH nor DNA-DNA
            SQL = """SELECT donor.amino, acceptor.amino FROM
                     contacts_%(pdb)s AS contacts
                        JOIN atoms_%(pdb)s AS donor
                           ON donor.id = contacts.donor_id
                        JOIN atoms_%(pdb)s AS acceptor
                           ON acceptor.id = contacts.acceptor_id
                     WHERE donor.amino = "HOH" AND acceptor.amino = "HOH"
                       OR ((donor.amino = "A" or
                                donor.amino = "G" or
                                donor.amino = "C" or
                                donor.amino = "U" or
                                donor.amino = "T") AND
                                (acceptor.amino = "A" or
                                acceptor.amino = "G" or
                                acceptor.amino = "C" or
                                acceptor.amino = "U" or
                                acceptor.amino = "T"))
                     """
            pdb = table.split('_')[1]
            r = cur2.execute(SQL % {'pdb' : pdb})
            r = r.fetchall()
            self.failUnless(len(r) == 0, "table %s:\n%s" % (table, r))
    
    def testC(self):
        cur2 = self.conn.cursor()
    	
    	atom_tables = cur2.execute("""SELECT name FROM sqlite_master
                    WHERE type='table' AND name LIKE 'atoms_%'
                    ORDER BY name""").fetchall()
        
        atom_table = iter(atom_tables)
        
        for (contacts_table,) in self.tables:
            (current, ) = atom_table.next()
            pdb_c = contacts_table.split('_')[1]
            pdb_a = current.split('_')[1]
            self.failUnless(pdb_c == pdb_a)
            
        self.assertRaises(StopIteration, atom_table.next) # there should be no excess atom tables
    
    def tearDown(self):
        self.conn.close()    
        
    def db_query(self, SQL, model_reply):
        self.cur.execute(SQL)
        for r in model_reply:
            self.failUnlessEqual(self.cur.fetchone(), r)

if __name__ == "__main__":
    DB = "db.sqlite"
    if len(sys.argv) > 1:
        DB = sys.argv.pop(1)
        
    unittest.main()
